import numpy as np
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 生成数据
X1 = 2 * np.random.rand(100, 1)
X2 = 2 * np.random.rand(100, 1)
Y = 4 + 3 * X1 + 5 * X2 + np.random.rand(100, 1)

# 有截距的模型
sk_linear = LinearRegression()
X_c = np.c_[X1, X2]
sk_linear.fit(X_c, Y)
print('有截距模型 - 截距:%s , 参数系数:%s' % (sk_linear.intercept_, sk_linear.coef_))

# 无截距的模型
lin_reg = LinearRegression(fit_intercept=False)
lin_reg.fit(X_c, Y)
print('无截距模型 - 参数系数:%s' % lin_reg.coef_)

# 创建预测网格
x1_range = np.linspace(0, 2, 20)
x2_range = np.linspace(0, 2, 20)
# 生成网格坐标
X1_grid, X2_grid = np.meshgrid(x1_range, x2_range)
X_grid = np.c_[X1_grid.ravel(), X2_grid.ravel()]

# 预测整个网格
Y_pred_sk = sk_linear.predict(X_grid).reshape(X1_grid.shape)
Y_pred_lin = lin_reg.predict(X_grid).reshape(X1_grid.shape)

# 绘制3D图形
fig = plt.figure(figsize=(12, 8)
                 , dpi=120,  # 更高分辨率
                 facecolor='#F5F5F5',  # 浅灰色背景
                 edgecolor='black',  # 黑色边框
                 frameon=True)  # 显示边框
ax = fig.add_subplot(111, projection='3d')
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
# 绘制原始数据点
ax.scatter(X1, X2, Y, c='b', marker='o', label='原始数据')

# 绘制有截距的回归平面
ax.plot_surface(X1_grid, X2_grid, Y_pred_sk,
                color='r', alpha=0.5, label='有截距模型')

# 绘制无截距的回归平面
ax.plot_surface(X1_grid, X2_grid, Y_pred_lin,
                color='g', alpha=0.3, label='无截距模型')

# 设置坐标轴标签
ax.set_xlabel('X1')
ax.set_ylabel('X2')
ax.set_zlabel('Y')
ax.set_title('3D线性回归比较')

# 添加图例
# 由于plot_surface不支持直接添加图例，我们需要手动创建代理对象
from matplotlib.patches import Patch

legend_elements = [
    Patch(facecolor='blue', alpha=0.5, label='原始数据'),
    Patch(facecolor='red', alpha=0.5, label='有截距模型'),
    Patch(facecolor='green', alpha=0.5, label='无截距模型')
]
ax.legend(handles=legend_elements)

plt.tight_layout()
plt.show()
